"""
数据加载和预处理模块
"""

import pandas as pd
import numpy as np
from typing import Dict, List, Tuple, Optional
import yaml
import os


class DataLoader:
    """数据加载器类"""
    
    def __init__(self, config_path: str = None):
        """
        初始化数据加载器
        
        Args:
            config_path: 配置文件路径
        """
        self.config = self._load_config(config_path)
        self.regions = self.config['hierarchy']['regions']
        
    def _load_config(self, config_path: str) -> Dict:
        """加载配置文件"""
        if config_path is None:
            config_path = "config/config.yaml"
        
        with open(config_path, 'r', encoding='utf-8') as f:
            config = yaml.safe_load(f)
        return config
    
    def process_dataframe(self, df: pd.DataFrame) -> Dict:
        """
        处理DataFrame数据
        
        Args:
            df: 已处理的时间序列DataFrame
            
        Returns:
            处理后的数据字典
        """
        print("处理DataFrame数据...")
        
        # 预处理数据
        processed_df = self.preprocess_dataframe(df)
        
        # 创建层级数据结构
        hierarchical_df = self.create_hierarchical_dataframe(processed_df)
        
        # 创建聚合数据
        aggregated_data = self.create_aggregated_data(hierarchical_df)
        
        # 计算比例数据
        proportions_data = self.calculate_proportions(aggregated_data)
        
        # 返回处理后的数据
        processed_data = {
            'raw': df,
            'processed': processed_df,
            'hierarchical': hierarchical_df,
            'aggregated': aggregated_data,
            'proportions': proportions_data
        }
        
        print("✓ DataFrame数据处理完成")
        return processed_data
    
    def preprocess_dataframe(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        预处理DataFrame数据
        
        Args:
            df: 原始DataFrame
            
        Returns:
            预处理后的DataFrame
        """
        # 处理缺失值
        df = df.fillna(0)
        
        # 确保数值列都是数值类型
        numeric_columns = df.select_dtypes(include=[np.number]).columns
        for col in numeric_columns:
            df[col] = pd.to_numeric(df[col], errors='coerce').fillna(0)
        
        # 处理时间列
        if '时间' in df.columns:
            df['时间'] = pd.to_datetime(df['时间'], format='%Y', errors='coerce')
            df['Year'] = df['时间'].dt.year
        else:
            # 如果没有时间列，从索引中提取
            df = df.reset_index()
            # 检查索引是否已经是年份格式
            if df.iloc[:, 0].dtype in ['int64', 'float64']:
                # 如果是数值，直接作为年份
                df['Year'] = df.iloc[:, 0].astype(int)
            else:
                # 如果是其他格式，尝试转换为年份
                try:
                    df['Year'] = pd.to_datetime(df.iloc[:, 0], format='%Y', errors='coerce').dt.year
                except:
                    # 如果转换失败，直接使用数值
                    df['Year'] = pd.to_numeric(df.iloc[:, 0], errors='coerce').fillna(0).astype(int)
        
        return df
    
    def create_hierarchical_dataframe(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        从DataFrame创建层级数据结构
        
        Args:
            df: 预处理后的DataFrame
            
        Returns:
            层级结构DataFrame
        """
        # 宽格式转长格式
        id_vars = ['Year']
        value_vars = [col for col in df.columns if col not in id_vars and col != '时间']
        
        df_long = df.melt(
            id_vars=id_vars,
            value_vars=value_vars,
            var_name='Province',
            value_name='Grain'
        )
        
        # 添加区域信息
        df_long['Area'] = df_long['Province'].apply(self._get_region)
        
        # 设置时间索引
        df_long['Time'] = pd.to_datetime(df_long['Year'], format='%Y')
        df_long = df_long.set_index('Time')
        
        return df_long
    
    def load_excel_data(self, file_path: str, sheet_names: List[str] = None) -> Dict[str, pd.DataFrame]:
        """
        加载Excel数据
        
        Args:
            file_path: Excel文件路径
            sheet_names: 工作表名称列表
            
        Returns:
            包含各工作表数据的字典
        """
        if sheet_names is None:
            sheet_names = self.config['data']['sheet_names']
            
        data_dict = {}
        for sheet_name in sheet_names:
            try:
                df = pd.read_excel(file_path, sheet_name=sheet_name)
                data_dict[sheet_name] = df
            except Exception as e:
                print(f"加载工作表 {sheet_name} 失败: {e}")
                
        return data_dict
    
    def preprocess_data(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        数据预处理
        
        Args:
            df: 原始数据框
            
        Returns:
            预处理后的数据框
        """
        # 处理缺失值
        df = df.fillna(0)
        
        # 确保年份列存在
        if 'Year' not in df.columns and '时间' in df.columns:
            df = df.rename(columns={'时间': 'Year'})
            
        # 确保数值列都是数值类型
        numeric_columns = df.select_dtypes(include=[np.number]).columns
        for col in numeric_columns:
            if col != 'Year':
                df[col] = pd.to_numeric(df[col], errors='coerce').fillna(0)
                
        return df
    
    def create_hierarchical_data(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        创建层级数据结构
        
        Args:
            df: 预处理后的数据框
            
        Returns:
            层级结构数据框
        """
        # 宽格式转长格式
        id_vars = ['Year'] if 'Year' in df.columns else []
        value_vars = [col for col in df.columns if col not in id_vars]
        
        df_long = df.melt(
            id_vars=id_vars,
            value_vars=value_vars,
            var_name='Province',
            value_name='Grain'
        )
        
        # 添加区域信息
        df_long['Area'] = df_long['Province'].apply(self._get_region)
        
        return df_long
    
    def _get_region(self, province: str) -> str:
        """根据省份获取所属区域"""
        for region, provinces in self.regions.items():
            if province in provinces:
                return region
        return "其他"
    
    def create_aggregated_data(self, df_long: pd.DataFrame) -> Dict[str, pd.DataFrame]:
        """
        创建聚合数据
        
        Args:
            df_long: 长格式数据框
            
        Returns:
            包含各层级聚合数据的字典
        """
        # 全国总量
        national_total = df_long.groupby('Year')['Grain'].sum().reset_index()
        national_total['Area'] = '全国'
        national_total['Province'] = '全国'
        
        # 区域总量
        regional_total = df_long.groupby(['Year', 'Area'])['Grain'].sum().reset_index()
        regional_total['Province'] = regional_total['Area']
        
        # 省份数据
        province_data = df_long.copy()
        
        return {
            'national': national_total,
            'regional': regional_total,
            'province': province_data
        }
    
    def calculate_proportions(self, aggregated_data: Dict[str, pd.DataFrame]) -> Dict[str, pd.DataFrame]:
        """
        计算比例数据
        
        Args:
            aggregated_data: 聚合数据字典
            
        Returns:
            包含比例数据的字典
        """
        national_total = aggregated_data['national']
        regional_total = aggregated_data['regional']
        province_data = aggregated_data['province']
        
        # 计算各省份在全国的占比
        province_national_prop = province_data.merge(
            national_total[['Year', 'Grain']], 
            on='Year', 
            suffixes=('', '_national')
        )
        province_national_prop['National_Proportion'] = (
            province_national_prop['Grain'] / province_national_prop['Grain_national']
        )
        
        # 计算各省份在区域中的占比
        province_regional_prop = province_data.merge(
            regional_total[['Year', 'Area', 'Grain']], 
            on=['Year', 'Area'], 
            suffixes=('', '_regional')
        )
        province_regional_prop['Regional_Proportion'] = (
            province_regional_prop['Grain'] / province_regional_prop['Grain_regional']
        )
        
        # 计算各区域在全国的占比
        regional_national_prop = regional_total.merge(
            national_total[['Year', 'Grain']], 
            on='Year', 
            suffixes=('', '_national')
        )
        regional_national_prop['National_Proportion'] = (
            regional_national_prop['Grain'] / regional_national_prop['Grain_national']
        )
        
        return {
            'province_national': province_national_prop,
            'province_regional': province_regional_prop,
            'regional_national': regional_national_prop
        }
    
    def load_data(self, file_path: str = None) -> Dict[str, pd.DataFrame]:
        """
        完整的数据加载流程
        
        Args:
            file_path: 数据文件路径
            
        Returns:
            处理后的数据字典
        """
        if file_path is None:
            file_path = self.config['data']['input_file']
            
        # 加载原始数据
        raw_data = self.load_excel_data(file_path)
        
        # 预处理数据
        processed_data = {}
        for sheet_name, df in raw_data.items():
            processed_df = self.preprocess_data(df)
            hierarchical_df = self.create_hierarchical_data(processed_df)
            aggregated_data = self.create_aggregated_data(hierarchical_df)
            proportions_data = self.calculate_proportions(aggregated_data)
            
            processed_data[sheet_name] = {
                'raw': df,
                'processed': processed_df,
                'hierarchical': hierarchical_df,
                'aggregated': aggregated_data,
                'proportions': proportions_data
            }
            
        return processed_data 